﻿// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the MIT License. See License.txt in the project root for license information.

using System;
using System.Collections.Generic;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Reactive.Linq;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.Azure.WebJobs.Host.Scale;
using Microsoft.Azure.WebJobs.Logging;
using Microsoft.Azure.WebJobs.Script.Diagnostics;
using Microsoft.Azure.WebJobs.Script.Eventing;
using Microsoft.Extensions.DependencyInjection;
using Microsoft.Extensions.Logging;
using Microsoft.Extensions.Options;
using Mono.Unix;

namespace Microsoft.Azure.WebJobs.Script.Workers
{
    internal abstract class WorkerProcess : IWorkerProcess, IDisposable
    {
        private readonly IProcessRegistry _processRegistry;
        private readonly ILogger _workerProcessLogger;
        private readonly IWorkerConsoleLogSource _consoleLogSource;
        private readonly IScriptEventManager _eventManager;
        private readonly IMetricsLogger _metricsLogger;
        private readonly int processExitTimeoutInMilliseconds = 1000;
        private readonly IServiceProvider _serviceProvider;
        private readonly IDisposable _eventSubscription;
        private readonly Lazy<ILogger> _toolingConsoleJsonLoggerLazy;
        private readonly IEnvironment _environment;
        private readonly IOptionsMonitor<ScriptApplicationHostOptions> _scriptApplicationHostOptions;

        private readonly object _syncLock = new();
        private readonly bool _useStdErrorStreamForErrorsOnly;
        private Queue<string> _processStdErrDataQueue = new(3);
        private IHostProcessMonitor _processMonitor;
        private TaskCompletionSource<int> _processExit; // used to hold custom exceptions on non-success exit.

        internal WorkerProcess(IScriptEventManager eventManager, IProcessRegistry processRegistry, ILogger workerProcessLogger, IWorkerConsoleLogSource consoleLogSource, IMetricsLogger metricsLogger,
            IServiceProvider serviceProvider, ILoggerFactory loggerFactory, IEnvironment environment, IOptionsMonitor<ScriptApplicationHostOptions> scriptApplicationHostOptions, bool useStdErrStreamForErrorsOnly = false)
        {
            _processRegistry = processRegistry;
            _workerProcessLogger = workerProcessLogger;
            _consoleLogSource = consoleLogSource;
            _eventManager = eventManager;
            _metricsLogger = metricsLogger;
            _useStdErrorStreamForErrorsOnly = useStdErrStreamForErrorsOnly;
            _serviceProvider = serviceProvider;
            _environment = environment;
            _scriptApplicationHostOptions = scriptApplicationHostOptions;
            _toolingConsoleJsonLoggerLazy = new Lazy<ILogger>(() => loggerFactory.CreateLogger(WorkerConstants.ToolingConsoleLogCategoryName), isThreadSafe: true);

            // We subscribe to host start events so we can handle the restart that occurs
            // on host specialization.
            _eventSubscription = _eventManager.OfType<HostStartEvent>().Subscribe(OnHostStart);
        }

        protected bool Disposing { get; private set; }

        public int Id => Process.Id;

        internal Queue<string> ProcessStdErrDataQueue => _processStdErrDataQueue;

        public Process Process { get; set; }

        public TaskCompletionSource<bool> ProcessWaitingForTermination { get; set; } = new TaskCompletionSource<bool>();

        internal abstract Process CreateWorkerProcess();

        public Task StartProcessAsync(CancellationToken cancellationToken = default)
        {
            cancellationToken.ThrowIfCancellationRequested();
            using (_metricsLogger.LatencyEvent(MetricEventNames.ProcessStart))
            {
                Process = CreateWorkerProcess();
                if (_environment.IsAnyLinuxConsumption())
                {
                    AssignUserExecutePermissionsIfNotExists();
                }

                _processExit = new();
                try
                {
                    Process.ErrorDataReceived += OnErrorDataReceived;
                    Process.OutputDataReceived += OnOutputDataReceived;
                    Process.Exited += OnProcessExited;
                    Process.EnableRaisingEvents = true;
                    string sanitizedArguments = Sanitizer.Sanitize(Process.StartInfo.Arguments);

                    _workerProcessLogger?.LogDebug($"Starting worker process with FileName:{Process.StartInfo.FileName} WorkingDirectory:{Process.StartInfo.WorkingDirectory} Arguments:{sanitizedArguments}");
                    Process.Start();
                    _workerProcessLogger?.LogDebug($"{Process.StartInfo.FileName} process with Id={Process.Id} started");

                    Process.BeginErrorReadLine();
                    Process.BeginOutputReadLine();

                    // Register process only after it starts
                    _processRegistry?.Register(this);

                    RegisterWithProcessMonitor();

                    return Task.CompletedTask;
                }
                catch (Exception ex)
                {
                    _processExit.TrySetException(ex);
                    _workerProcessLogger.LogError(ex, $"Failed to start Worker Channel. Process fileName: {Process.StartInfo.FileName}");
                    return Task.FromException(ex);
                }
            }
        }

        public Task<int> WaitForExitAsync(CancellationToken cancellationToken = default)
        {
            ObjectDisposedException.ThrowIf(Disposing, this);
            if (_processExit is { } tcs)
            {
                // We use a TaskCompletionSource (and not Process.WaitForExitAsync) so we can propagate our custom exceptions.
                return tcs.Task.WaitAsync(cancellationToken);
            }

            throw new InvalidOperationException("Process has not been started yet.");
        }

        private void OnErrorDataReceived(object sender, DataReceivedEventArgs e)
        {
            if (e.Data != null)
            {
                ParseErrorMessageAndLog(e.Data);
            }
        }

        internal void ParseErrorMessageAndLog(string msg)
        {
            if (msg.IndexOf("warn", StringComparison.OrdinalIgnoreCase) > -1)
            {
                BuildAndLogConsoleLog(msg, LogLevel.Warning);
            }
            else
            {
                if (_useStdErrorStreamForErrorsOnly)
                {
                    LogError(msg);
                }
                else
                {
                    // TODO: redesign how we log errors so it's not based on the string contents (GH issue #8273)
                    if ((msg.IndexOf("error", StringComparison.OrdinalIgnoreCase) > -1) ||
                        (msg.IndexOf("fail", StringComparison.OrdinalIgnoreCase) > -1) ||
                        (msg.IndexOf("severe", StringComparison.OrdinalIgnoreCase) > -1) ||
                        (msg.IndexOf("unhandled exception", StringComparison.OrdinalIgnoreCase) > -1))
                    {
                        LogError(msg);
                    }
                    else
                    {
                        BuildAndLogConsoleLog(msg, LogLevel.Information);
                    }
                }
            }
        }

        private void LogError(string msg)
        {
            BuildAndLogConsoleLog(msg, LogLevel.Error);
            _processStdErrDataQueue = WorkerProcessUtilities.AddStdErrMessage(_processStdErrDataQueue, Sanitizer.Sanitize(msg));
        }

        private void OnProcessExited(object sender, EventArgs e)
        {
            _workerProcessLogger.LogDebug("Process {processId} has exited with code {exitCode}.", Process?.Id, Process?.ExitCode);

            if (Disposing)
            {
                return;
            }

            int exit = 0;
            try
            {
                ThrowIfExitError();

                exit = Process.ExitCode;
                if (Process.ExitCode == WorkerConstants.SuccessExitCode)
                {
                    Process.WaitForExit();
                    Process.Close();
                }
                else if (Process.ExitCode == WorkerConstants.IntentionalRestartExitCode)
                {
                    HandleWorkerProcessRestart();
                }
            }
            catch (WorkerProcessExitException processExitEx)
            {
                _processExit.TrySetException(processExitEx);
                HandleWorkerProcessExitError(processExitEx);
            }
            catch (Exception exc)
            {
                // ignore process is already disposed
                _processExit.TrySetException(exc);
                _workerProcessLogger?.LogDebug(exc, "Exception on worker process exit. Process id: {processId}", Process?.Id);
            }
            finally
            {
                _processExit.TrySetResult(exit);
                UnregisterFromProcessMonitor();
            }
        }

        private void ThrowIfExitError()
        {
            if (Process.ExitCode is WorkerConstants.SuccessExitCode or WorkerConstants.IntentionalRestartExitCode)
            {
                return;
            }

            string exceptionMessage = string.Join(",", _processStdErrDataQueue.Where(s => !string.IsNullOrEmpty(s)));
            string sanitizedExceptionMessage = Sanitizer.Sanitize(exceptionMessage);
            WorkerProcessExitException processExitEx = new(
                $"{Process.StartInfo.FileName} exited with code {Process.ExitCode} (0x{Process.ExitCode:X})",
                new Exception(sanitizedExceptionMessage))
            {
                ExitCode = Process.ExitCode,
                Pid = Process.Id
            };

            throw processExitEx;
        }

        private void OnOutputDataReceived(object sender, DataReceivedEventArgs e)
        {
            if (e.Data != null)
            {
                BuildAndLogConsoleLog(e.Data, LogLevel.Information);
            }
        }

        internal void BuildAndLogConsoleLog(string msg, LogLevel level)
        {
            ConsoleLog consoleLog = new ConsoleLog()
            {
                Message = msg,
                Level = level
            };

            if (WorkerProcessUtilities.IsConsoleLog(msg))
            {
                _workerProcessLogger?.Log(level, WorkerProcessUtilities.RemoveLogPrefix(msg));
            }
            else if (WorkerProcessUtilities.IsToolingConsoleJsonLogEntry(msg))
            {
                // log with the message prefix as coretools expects it.
                _toolingConsoleJsonLoggerLazy.Value.Log(level, msg);
            }
            else
            {
                // These are 'user' console logs and should be forwarded to the ScriptHost, unless
                // we're in placeholder mode. In that case, we ignore the logs so they are not captured
                // by the specialized host when it starts.
                if (!_scriptApplicationHostOptions.CurrentValue.IsStandbyConfiguration)
                {
                    _consoleLogSource?.Log(consoleLog);
                }
            }
        }

        internal abstract void HandleWorkerProcessExitError(WorkerProcessExitException langExc);

        internal abstract void HandleWorkerProcessRestart();

        public void WaitForProcessExitInMilliSeconds(int waitTime)
        {
            try
            {
                if (!Process.HasExited)
                {
                    Process.WaitForExit(waitTime);
                }
            }
            catch (Exception ex)
            {
                _workerProcessLogger.LogDebug(ex, "An exception was thrown while waiting for a worker process to exit. It is possible that the process had already exited and this can be ignored.");
            }
        }

        public void Dispose()
        {
            Disposing = true;
            // best effort process disposal

            ProcessWaitingForTermination.SetResult(false);

            try
            {
                _eventSubscription?.Dispose();

                if (Process != null)
                {
                    if (!Process.HasExited)
                    {
                        Process.Kill();
                        if (!Process.WaitForExit(processExitTimeoutInMilliseconds))
                        {
                            _workerProcessLogger.LogWarning("Worker process {processId} has not exited despite waiting for {processExitTimeoutInMilliseconds} ms", Process?.Id, processExitTimeoutInMilliseconds);
                        }
                    }
                    Process.Dispose();
                }
            }
            catch (Exception exc)
            {
                _workerProcessLogger?.LogDebug(exc, "Exception on worker disposal.");
                //ignore
            }
        }

        internal void OnHostStart(HostStartEvent evt)
        {
            if (!Disposing)
            {
                RegisterWithProcessMonitor();
            }
        }

        /// <summary>
        /// Ensures that our process is registered with <see cref="IHostProcessMonitor"/>.
        /// </summary>
        /// <remarks>
        /// The goal is to ensure that all worker processes are registered with the monitor for the active host.
        /// There are a few different cases to consider:
        /// - Starting up in normal mode, we register on when the process is started.
        /// - When a placeholder mode host is specialized, a new host will be started but the previously initialized
        ///   worker process will remain running. We need to re-register with the new host.
        /// - When the worker process dies and a new instance of this class is created.
        /// </remarks>
        internal void RegisterWithProcessMonitor()
        {
            var processMonitor = _serviceProvider.GetScriptHostServiceOrNull<IHostProcessMonitor>();
            lock (_syncLock)
            {
                if (processMonitor != null && processMonitor != _processMonitor && Process != null)
                {
                    processMonitor.RegisterChildProcess(Process);
                    _processMonitor = processMonitor;
                }
            }
        }

        internal void UnregisterFromProcessMonitor()
        {
            lock (_syncLock)
            {
                if (_processMonitor != null && Process != null)
                {
                    // if we've registered our process with the monitor, unregister
                    _processMonitor.UnregisterChildProcess(Process);
                    _processMonitor = null;
                }
            }
        }

        private void AssignUserExecutePermissionsIfNotExists()
        {
            try
            {
                if (Process is not { } p)
                {
                    return;
                }

                string filePath = p.StartInfo.FileName;
                if (!File.Exists(filePath))
                {
                    _workerProcessLogger.LogDebug("File path does not exist, not assigning permissions: {filePath}", filePath);
                    return;
                }

                UnixFileInfo fileInfo = new(filePath);
                if (!fileInfo.FileAccessPermissions.HasFlag(FileAccessPermissions.UserExecute))
                {
                    _workerProcessLogger.LogDebug("Assigning execute permissions to file: {filePath}", filePath);
                    fileInfo.FileAccessPermissions |= FileAccessPermissions.UserExecute |
                                                      FileAccessPermissions.GroupExecute |
                                                      FileAccessPermissions.OtherExecute;
                }
            }
            catch (Exception ex)
            {
                _workerProcessLogger.LogWarning(ex, "Error while assigning execute permission.");
            }
        }
    }
}
